import cv2
import time
import numpy as np
import json
from ultralytics import YOLO
from collections import defaultdict
import os
from app.model.AccidentDetector import AccidentDetector
from collections import deque
from app.model.CollisionDetector import CollisionDetector


class YoloModel:
    def __init__(self, model_path,
                 traffic_flow: bool,
                 num_lanes: int,
                 max_trail_length=30,
                 calibration_file='../calibration/lane_calibration.json',
                 road_curvature=0.3,
                 accident_detection=True,
                 accident_params=None,
                 accident_types=None):
        self.model = YOLO(model_path)
        self.num_lanes = num_lanes
        self.max_trail_length = max_trail_length
        self.road_curvature = road_curvature
        self.accident_detection = accident_detection

        # 固定鸟瞰图尺寸
        self.bird_view_width = 1980
        self.bird_view_height = 1020

        # 车道线点数据
        self.left_lane_points = []
        self.right_lane_points = []
        self.center_lane_points = []
        self.lane_polygon = None  # 用于鸟瞰图显示的车道区域

        # 加载车道线校准数据
        self.load_calibration_data(calibration_file)

        # 构建车道区域多边形（用于鸟瞰图显示）
        self.build_lane_polygon_for_display()

        # 跟踪历史
        self.track_history = defaultdict(list)

        # 车辆计数
        self.vehicle_count = 0
        self.category_count = defaultdict(int)
        self.seen_track_ids = set()

        # 交通流计数
        self.traffic_flow = traffic_flow
        self.crossed_ids = set()
        self.crossing_count = 0

        # 事故检测器列表
        self.accident_detectors = []

        # 初始化车祸检测器
        if self.accident_detection:
            params = accident_params or {}
            debug_enabled = params.get('debug_enabled', False)

            # 默认启用所有检测器
            if accident_types is None:
                accident_types = ['off_road', 'collision']

            # 创建离开道路检测器
            if 'off_road' in accident_types:
                off_road_detector = AccidentDetector(
                    bird_view_width=self.bird_view_width,
                    bird_view_height=self.bird_view_height,
                    debug_enabled=debug_enabled
                )
                self.accident_detectors.append(off_road_detector)

            # 创建碰撞检测器
            if 'collision' in accident_types:
                collision_detector = CollisionDetector(
                    bird_view_width=self.bird_view_width,
                    bird_view_height=self.bird_view_height,
                    debug_enabled=debug_enabled
                )
                self.accident_detectors.append(collision_detector)

            # 为需要车道信息的检测器构建车道多边形
            for detector in self.accident_detectors:
                if hasattr(detector, 'build_lane_polygon') and self.left_lane_points and self.right_lane_points:
                    # 使用原始点构建检测多边形（不排序）
                    detector.build_lane_polygon(
                        self.left_lane_points,
                        self.right_lane_points
                    )

    def build_lane_polygon_for_display(self):
        """构建用于显示的车道区域多边形（严格保持之前正确版本）"""
        if not self.left_lane_points or not self.right_lane_points:
            return None

        # 确保左右车道点数量相同
        n_points = min(len(self.left_lane_points), len(self.right_lane_points))

        # 找到最底部的点（y坐标最大的点）
        max_y_left = max(p[1] for p in self.left_lane_points)
        max_y_right = max(p[1] for p in self.right_lane_points)
        max_y = max(max_y_left, max_y_right)

        # 如果最底部的点没有到达图像底部，则添加垂直延伸
        if max_y < self.bird_view_height - 1:
            # 获取左右车道最底部的点
            left_bottom = next(p for p in self.left_lane_points if p[1] == max_y_left)
            right_bottom = next(p for p in self.right_lane_points if p[1] == max_y_right)

            # 创建垂直延伸的点
            left_extension = (left_bottom[0], self.bird_view_height - 1)
            right_extension = (right_bottom[0], self.bird_view_height - 1)
        else:
            # 如果已经到达底部，则不需要添加延伸
            left_extension = None
            right_extension = None

        # 创建车道区域多边形
        polygon_points = []

        # 添加左车道线点（按原始顺序）
        for i in range(n_points):
            polygon_points.append(self.left_lane_points[i])

        # 如果添加了左车道延伸点
        if left_extension:
            polygon_points.append(left_extension)

        # 添加右车道延伸点（如果有）
        if right_extension:
            polygon_points.append(right_extension)

        # 添加右车道线点（按原始顺序，反向）
        for i in range(n_points - 1, -1, -1):
            polygon_points.append(self.right_lane_points[i])

        # 转换为NumPy数组
        self.lane_polygon = np.array(polygon_points, dtype=np.int32)

    def load_calibration_data(self, calibration_file):
        """加载车道线校准数据"""
        if not calibration_file:
            raise ValueError("Calibration file path is required")
        if not os.path.exists(calibration_file):
            raise FileNotFoundError(f"Calibration file not found: {calibration_file}")

        try:
            with open(calibration_file, 'r') as f:
                calibration_data = json.load(f)

            # 提取车道线点
            self.left_lane_points = calibration_data.get('left_lane_points', [])
            self.right_lane_points = calibration_data.get('right_lane_points', [])
            self.center_lane_points = calibration_data.get('center_lane_points', [])

            print(f"成功加载车道线校准数据: {calibration_file}")
            print(f"左车道线点数: {len(self.left_lane_points)}")
            print(f"右车道线点数: {len(self.right_lane_points)}")
            print(f"中心车道线点数: {len(self.center_lane_points)}")

        except Exception as e:
            raise RuntimeError(f"加载校准数据失败: {str(e)}")

    @staticmethod
    def draw_dashed_line(img, start, end, color, thickness, dash_length, gap_length):
        """绘制虚线"""
        total_length = int(np.sqrt((end[0] - start[0]) ** 2 + (end[1] - start[1]) ** 2))
        num_dashes = total_length // (dash_length + gap_length)

        for i in range(num_dashes):
            dash_start = (
                int(start[0] + (end[0] - start[0]) * (i * (dash_length + gap_length)) / total_length),
                int(start[1] + (end[1] - start[1]) * (i * (dash_length + gap_length)) / total_length)
            )
            dash_end = (
                int(start[0] + (end[0] - start[0]) * ((i * (dash_length + gap_length)) + dash_length) / total_length),
                int(start[1] + (end[1] - start[1]) * ((i * (dash_length + gap_length)) + dash_length) / total_length)
            )
            cv2.line(img, dash_start, dash_end, color, thickness)

    def draw_lane_lines(self, birdView_frame):
        """绘制车道线（严格保持之前正确版本）"""
        # 使用加载的车道线点
        left_points = self.left_lane_points
        right_points = self.right_lane_points
        center_points = self.center_lane_points

        # 1. 绘制左侧车道线 (绿色)
        if left_points:
            pts = np.array(left_points, np.int32).reshape((-1, 1, 2))
            cv2.polylines(birdView_frame, [pts], False, (0, 255, 0), 2, lineType=cv2.LINE_AA)

        # 2. 绘制右侧车道线 (红色)
        if right_points:
            pts = np.array(right_points, np.int32).reshape((-1, 1, 2))
            cv2.polylines(birdView_frame, [pts], False, (0, 0, 255), 2, lineType=cv2.LINE_AA)

        # 3. 绘制中心虚线 (黄色)
        if center_points:
            for i in range(0, len(center_points) - 1, 2):
                start = center_points[i]
                end = center_points[i + 1]
                cv2.line(birdView_frame,
                         (int(start[0]), int(start[1])),
                         (int(end[0]), int(end[1])),
                         (0, 255, 255), 2, lineType=cv2.LINE_AA)

        # 绘制车道区域（半透明）- 使用专门构建的显示多边形
        if self.lane_polygon is not None:
            overlay = birdView_frame.copy()
            cv2.fillPoly(overlay, [self.lane_polygon], (0, 100, 0))
            cv2.addWeighted(overlay, 0.2, birdView_frame, 0.8, 0, birdView_frame)

    def draw_trajectories(self, birdView_frame, frame_shape):
        """在鸟瞰图上绘制车辆运动轨迹"""
        frame_height, frame_width = frame_shape[:2]

        # 计算缩放比例
        scale_x = self.bird_view_width / frame_width
        scale_y = self.bird_view_height / frame_height

        # 为每个轨迹分配固定颜色
        colors = [
            (255, 0, 0), (0, 255, 0), (0, 0, 255),
            (255, 255, 0), (255, 0, 255), (0, 255, 255), (255, 255, 255)
        ]

        # 绘制所有轨迹
        for track_id, track in self.track_history.items():
            if len(track) < 2:
                continue

            # 获取轨迹颜色
            color_idx = track_id % len(colors)
            color = colors[color_idx]

            # 缩放轨迹点以适应鸟瞰图
            bird_track = []
            for point in track:
                scaled_x = int(point[0] * scale_x)
                scaled_y = int(point[1] * scale_y)
                bird_track.append((scaled_x, scaled_y))

            # 绘制完整轨迹线
            for i in range(1, len(bird_track)):
                cv2.line(birdView_frame, bird_track[i - 1], bird_track[i], color, 3)

            # 在轨迹起点绘制标记
            if bird_track:
                cv2.circle(birdView_frame, bird_track[0], 8, color, -1)

            if bird_track:
                if len(bird_track) > 2:
                    dx = bird_track[-1][0] - bird_track[-2][0]
                    dy = bird_track[-1][1] - bird_track[-2][1]
                    arrow_length = 20
                    angle = np.arctan2(dy, dx)
                    pt1 = (int(bird_track[-1][0] + arrow_length * np.cos(angle + np.pi / 6)),
                           int(bird_track[-1][1] + arrow_length * np.sin(angle + np.pi / 6)))
                    pt2 = (int(bird_track[-1][0] + arrow_length * np.cos(angle - np.pi / 6)),
                           int(bird_track[-1][1] + arrow_length * np.sin(angle - np.pi / 6)))
                    pts = np.array([bird_track[-1], pt1, pt2], np.int32)
                    cv2.fillPoly(birdView_frame, [pts], color)
                else:
                    cv2.circle(birdView_frame, bird_track[-1], 10, color, -1)

                # 显示车辆ID
                cv2.putText(birdView_frame, f"ID:{track_id}",
                            (bird_track[-1][0] + 15, bird_track[-1][1]),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)

    def track(self, frame):
        """执行目标跟踪和注释（严格保持之前正确版本）"""
        if frame is None or frame.size == 0:
            raise ValueError("Invalid input frame: frame is None or empty")

        try:
            # 执行目标跟踪
            results = self.model.track(frame, persist=True, show=False, verbose=False)
        except Exception as e:
            print(f"Tracking error: {e}")
            birdView_frame = np.zeros((self.bird_view_height, self.bird_view_width, 3), dtype=np.uint8)
            self.draw_lane_lines(birdView_frame)
            return frame, frame, birdView_frame

        boxes = results[0].boxes.xywh.cpu()
        track_ids = results[0].boxes.id.int().cpu().tolist() if results[0].boxes.id is not None else []
        cls_indices = results[0].boxes.cls.int().cpu().tolist()
        class_names = self.model.names

        # 创建带注释的帧和鸟瞰图
        annotated_frame = results[0].plot()
        birdView_frame = np.zeros((self.bird_view_height, self.bird_view_width, 3), dtype=np.uint8)

        # 绘制车道线
        self.draw_lane_lines(birdView_frame)

        # 添加位置标识和时间戳
        cv2.putText(birdView_frame, "G243-K992+840", (150, self.bird_view_height - 50),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 200, 255), 1, cv2.LINE_AA)
        current_time = time.strftime("%Y-%m-%d %A %H:%M:%S", time.localtime())
        cv2.putText(birdView_frame, current_time, (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (100, 255, 255), 1, cv2.LINE_AA)

        # 获取帧尺寸
        frame_height, frame_width = frame.shape[:2]

        # 处理每个检测到的目标
        for box, track_id, cls_idx in zip(boxes, track_ids, cls_indices):
            class_name = class_names[cls_idx]

            # 新车统计
            if track_id not in self.seen_track_ids:
                self.seen_track_ids.add(track_id)
                self.vehicle_count += 1
                self.category_count[class_name] += 1

            # 获取目标中心点
            x, y, w, h = box
            track = self.track_history[track_id]

            # 存储车辆中心点（原始帧坐标）
            center_x = float(x)
            center_y = float(y)
            track.append((center_x, center_y))

            # 限制轨迹长度
            if len(track) > self.max_trail_length:
                track.pop(0)

            # 在原始帧上绘制轨迹
            points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(annotated_frame, [points], isClosed=False, color=(230, 230, 230), thickness=10)

            # 交通流统计
            if self.traffic_flow:
                frame_height = frame.shape[0]
                middle_line_y = frame_height // 2

                # 绘制黄色虚线中线
                self.draw_dashed_line(annotated_frame, (0, middle_line_y), (frame.shape[1], middle_line_y),
                                      (0, 255, 255), 2, 20, 10)

                # 检测是否穿越中线
                if len(track) >= 2:
                    prev_y = track[-2][1]
                    curr_y = center_y
                    if (prev_y < middle_line_y <= curr_y) or (prev_y > middle_line_y >= curr_y):
                        if track_id not in self.crossed_ids:
                            self.crossed_ids.add(track_id)
                            self.crossing_count += 1

        # 在鸟瞰图上绘制所有车辆轨迹
        self.draw_trajectories(birdView_frame, frame.shape)

        # 事故检测（多个检测器）
        all_accidents = []
        if self.accident_detection and self.accident_detectors:
            for detector in self.accident_detectors:
                annotated_frame, accidents = detector.detect_accidents(
                    annotated_frame,
                    self.track_history,
                    frame_width,
                    frame_height
                )
                all_accidents.extend(accidents)

        # 在鸟瞰图上添加事故计数
        if all_accidents:
            cv2.putText(birdView_frame, f"ACCIDENTS: {len(all_accidents)}",
                        (self.bird_view_width - 300, 30),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 255), 2)

            # 在processed视图中添加事故类型标记
            for accident in all_accidents:
                # 事故数据格式: (描述, 类型)
                accident_type = accident[1] if len(accident) > 1 else "事故"

                # 在processed视图左上角添加标记
                cv2.putText(annotated_frame, f"{accident_type} DETECTED!",
                            (20, 40),
                            cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)

                # 在processed视图右上角添加总事故计数
                cv2.putText(annotated_frame, f"Total Accidents: {len(all_accidents)}",
                            (annotated_frame.shape[1] - 300, 40),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)

        return annotated_frame, frame, birdView_frame

    def get_statistics(self):
        """获取当前统计信息"""
        stats = {
            "total_count": self.vehicle_count,
            "category_count": dict(self.category_count),
            "crossing_count": self.crossing_count,
            "current_tracks": len(self.track_history),
            "accident_count": 0,
            "accident_details": []
        }

        # 收集所有检测器的事故信息
        if self.accident_detectors:
            all_accidents = []
            for detector in self.accident_detectors:
                accidents = detector.get_accident_history()
                if accidents:
                    all_accidents.extend(accidents)

            stats["accident_count"] = len(all_accidents)
            stats["accident_details"] = all_accidents

        return stats

    def reset_statistics(self):
        """重置所有统计计数器和跟踪器"""
        self.vehicle_count = 0
        self.category_count = defaultdict(int)
        self.seen_track_ids = set()
        self.crossing_count = 0
        self.crossed_ids.clear()
        self.track_history.clear()

        # 重置事故检测器状态
        for detector in self.accident_detectors:
            if hasattr(detector, 'reset'):
                detector.reset()